package com.lucifer.hawkeye.ai.rag;
import com.lucifer.hawkeye.ai.domain.rag.type.RagType;
import com.lucifer.hawkeye.ai.vector.VectorStoreEngine;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.lang.Nullable;

import java.util.Comparator;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @author lucifer
 * @date 2025/3/28 17:46
 */
@Slf4j
public class HawkeyeRagEngine implements RagEngine {

    private final VectorStoreEngine vectorStoreEngine;

    @Override
    public boolean add(Document document) {
        try {
            synchronized (vectorStoreEngine) {
                vectorStoreEngine.add(List.of(document));
                return true;
            }
        } catch (Exception e) {
            log.error("document id: {} added Failed: {}", document.getId(), e);
            return false;
        }
    }

    @Override
    public List<Document> similaritySearch(String conversationId, String ragType, String content, String tableName) {
        FilterExpressionBuilder expression = new FilterExpressionBuilder();
        Filter.Expression filterExpression = null;
        SearchRequest.Builder searchRequestBuilder = SearchRequest.builder();
        if (Objects.equals(RagType.DDL.name(), ragType)) {
            filterExpression = expression.and(expression.eq("rag_type", ragType), expression.eq("table_name", tableName)).build();
            searchRequestBuilder.query(tableName);
        }
        if (Objects.equals(RagType.SQL.name(), ragType)) {
            filterExpression = expression.and(expression.eq("rag_type", ragType), expression.eq("semantic", content)).build();
            //filterExpression = expression.eq("rag_type", ragType).build();
            searchRequestBuilder.query(content);
        }
        SearchRequest searchRequest = searchRequestBuilder
                .topK(10)
                .filterExpression(filterExpression)
                .build();
        List<Document> documentList = vectorStoreEngine.similaritySearch(searchRequest);
        List<Document> sortedList = documentList.stream()
                .filter(doc -> doc.getScore() >= 0.4)
                .sorted(Comparator.comparing(doc -> doc.getScore(), Comparator.reverseOrder()))
                .collect(Collectors.toList());

        List<Document> list = Lists.newArrayList();
        if(!sortedList.isEmpty()){
            list.add(sortedList.get(0));
        }
        return list;
    }


    public HawkeyeRagEngine(VectorStoreEngine vectorStoreEngine) {
        this.vectorStoreEngine = vectorStoreEngine;
    }

    public static Builder builder() {
        return new Builder();
    }


    public static final class Builder {

        @Nullable
        private VectorStoreEngine vectorStoreEngine;

        public Builder() {
        }

        public Builder vectorStoreEngine(VectorStoreEngine vectorStoreEngine) {
            this.vectorStoreEngine = vectorStoreEngine;
            return this;
        }

        public HawkeyeRagEngine build() {
            return new HawkeyeRagEngine(this.vectorStoreEngine);
        }
    }
}
